import torch
import torch.nn as nn
import torch.nn.functional as F

from .bconv import Bconv
from .cfm import CFM
from .fam import FAM
from .pdc import PDC
from .resnet import ResNet18,init_weight
from .rf import RF
from .sa import SA
from .sca import SpatialAttention, ChannelwiseAttention

class MyNet_base(nn.Module):
    def __init__(self):
        super(MyNet_base, self).__init__()
        self.resnet = ResNet18()
        init_weight(self.resnet)
        self.downsample2 = nn.MaxPool2d(2, stride=2)
        self.upsameple2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv1 = Bconv(512,256,3,1,1)
        self.conv2 = Bconv(256,128,3,1,1)
        self.conv3 = Bconv(128,64,3,1,1)
        self.conv4 = Bconv(64,1,3,1,1)

    def forward(self,x):
        resnet_dict = self.resnet(x)
        x0 = resnet_dict['X0']      # 64 88 88
        x1 = resnet_dict['X1']      # 64 88 88
        x2 = resnet_dict['X2']      # 128 44 44
        x3 = resnet_dict['X3']      # 256 22 22
        x4 = resnet_dict['X4']      # 512 11 11

        up_x4 = self.upsameple2(x4)
        out4 = up_x4

        out3 = self.conv1(out4)
        out3 = out3 + x3
        out3 = self.upsameple2(out3)

        out2 = self.conv2(out3)
        out2 = out2 + x2
        out2 = self.upsameple2(out2)

        out1 = self.conv3(out2)
        out1 = out1 + x1

        out = self.conv4(out1)
        out = self.upsameple2(self.upsameple2(out))

        return out, out1, out2, out3, out4